import csv
import argparse
import math
# Helper function to generate server-to-leaf connections based on the "NVIDIA style"
# where each server connects to multiple (e.g., 8) different leaf switches.
# This pattern is derived from the provided instruction and resembles common NVIDIA GPU cluster topologies.
def _generate_server_to_leaf_connections_nvidia_style(
    num_servers,
    num_leaf_switches,
    ports_per_leaf,  # Needed to determine downlink port assignment on leaf
    server_template,
    leaf_template,
    server_ports_to_use=8, # Number of ports on server to connect to leaves
    server_port_prefix="eth", # Prefix for server port names, e.g., eth1, eth2
    leaf_port_prefix="Ethernet" # Prefix for leaf port names, e.g., Ethernet1
):
    """
    Generates server-to-leaf connections following a specific multi-rail pattern.
    Each server uses a defined number of its ports (e.g., 8) to connect to
    an equal number of distinct leaf switches.

    The connection pattern is:
    - Server S uses its ports P1, P2, ..., Pn (e.g., eth1, eth2, ..., eth8).
    - Port Pi of Server S connects to Leaf L_i.
    - The target Leaf L_i is determined by a round-robin assignment involving both
      the server ID and the server's port index. Specifically, for server S (1-based ID)
      and its j-th port (0-based index, for P(j+1)), the target leaf ID is
      `((S-1) + j) % num_leaf_switches + 1`.
    - All connections from Server S to its respective leaves use the same port number
      on the leaf side. This leaf port number is determined by the server ID, cycling
      through the available downlink ports on a leaf: `((S-1) % downlinks_per_leaf) + 1`.

    Args:
        num_servers (int): Total number of servers.
        num_leaf_switches (int): Total number of leaf switches. Caller must ensure
                                 this is sufficient, ideally >= server_ports_to_use.
        ports_per_leaf (int): Number of ports on each leaf switch.
        server_template (str): Naming template for servers (e.g., "gpu{}").
        leaf_template (str): Naming template for leaf switches (e.g., "leaf{}").
        server_ports_to_use (int): Number of ports on each server dedicated to
                                   leaf connections. Defaults to 8.
        server_port_prefix (str): Prefix for server port names (e.g., "eth").
        leaf_port_prefix (str): Prefix for leaf port names (e.g., "Ethernet").

    Returns:
        list: A list of tuples, where each tuple is (adev, aport, zdev, zport).
    
    Raises:
        ValueError: If ports_per_leaf is too small to allow for downlink ports.
    """
    connections = []
    if num_servers == 0:
        return connections

    # It's crucial that num_leaf_switches is adequate.
    # If num_leaf_switches < server_ports_to_use, a single server cannot connect to
    # distinct leaves for all its specified ports; leaf IDs will wrap around.
    # This should ideally be handled by the caller's logic for calculating num_leaf_switches.
    if num_leaf_switches < server_ports_to_use:
        print(
            f"Warning: num_leaf_switches ({num_leaf_switches}) is less than "
            f"server_ports_to_use ({server_ports_to_use}). "
            "A server might connect to the same leaf multiple times via different ports."
        )

    downlinks_per_leaf = ports_per_leaf // 2
    if downlinks_per_leaf <= 0:
        raise ValueError(
            f"ports_per_leaf ({ports_per_leaf}) must be configured to allow for "
            "at least one downlink port (i.e., ports_per_leaf >= 2)."
        )

    for server_idx_0based in range(num_servers):
        server_id_1based = server_idx_0based + 1
        server_name = server_template.format(server_id_1based)

        # Determine the group the server belongs to (0-based group index)
        server_group_idx_0based = server_idx_0based // 32
        # Calculate the base leaf ID for this group (1-based global ID)
        base_leaf_id_for_group = server_group_idx_0based * 8 + 1

        # Determine the port number on the leaf switch that this server's connections will use.
        # This port number is consistent for all connections originating from this server.
        # Example: Server 1 uses port 1 on its target leaves, Server 2 uses port 2, etc.,
        # cycling based on the number of available downlink ports on a leaf.
        # This still uses the server's global index to distribute load across downlink ports.
        leaf_port_number_on_target_leaf = (server_idx_0based % downlinks_per_leaf) + 1
        leaf_zport_name = f"{leaf_port_prefix}{leaf_port_number_on_target_leaf}"

        for server_port_idx_0based in range(server_ports_to_use):
            # Server's physical port number (1-based, e.g., 1 through 8)
            server_physical_port_num = server_port_idx_0based + 1
            server_aport_name = f"{server_port_prefix}{server_physical_port_num}"

            # Determine the target leaf switch for this specific server port within the group's leaves.
            # Each of the server_ports_to_use connects to a different leaf within the 8 leaves of the group.
            # The leaf ID within the group is determined by the server port index (round-robin 0-7).
            leaf_id_within_group_0based = server_port_idx_0based % 8
            target_leaf_id_1based = base_leaf_id_for_group + leaf_id_within_group_0based

            # Ensure the target leaf ID does not exceed the total number of leaf switches available.
            # This check might be redundant if num_leaf_switches is calculated correctly in generate_connections,
            # but it adds robustness.
            if target_leaf_id_1based > num_leaf_switches:
                 print(f"Warning: Calculated target leaf ID {target_leaf_id_1based} exceeds total available leaf switches {num_leaf_switches}. Adjusting to the last available leaf.")
                 target_leaf_id_1based = num_leaf_switches # Or raise an error, depending on desired behavior for insufficient leaves

            target_leaf_name = leaf_template.format(target_leaf_id_1based)

            connections.append((
                server_name,
                server_aport_name,
                target_leaf_name,
                leaf_zport_name
            ))
            
    return connections

def generate_connections(
    num_servers,
    ports_per_leaf,
    ports_per_spine,
    server_template="server{}",
    leaf_template="leaf{}",
    spine_template="spine{}",
    # core_template="core{}", # Parameter available if 3-level topology is added later
    # ports_per_core=0,       # Parameter available if 3-level topology is added later
    output_filename="cables.csv",
    server_start_port=1, # This parameter is no longer directly used for NVIDIA style connections, but keeping for potential compatibility or future use.
    server_ports_to_use=8, # New parameter for NVIDIA style
    server_port_prefix="eth", # New parameter for NVIDIA style
    leaf_port_prefix="Ethernet" # New parameter for NVIDIA style
):
    """
    Generates a list of network connections for a 2-level (Leaf-Spine) fat-tree topology.
    This topology is often referred to as "rail-optimized" in contexts like NVIDIA designs.

    Args:
        num_servers (int): Total number of servers.
        ports_per_leaf (int): Number of ports per leaf switch. Must be even for the
                              specified 50/50 downlink/uplink port split.
        ports_per_spine (int): Number of ports per spine switch.
        server_template (str): Naming template for servers (e.g., "gpu{}").
                               The placeholder {} will be filled with a 1-based ID.
        leaf_template (str): Naming template for leaf switches.
                             The placeholder {} will be filled with a 1-based ID.
        spine_template (str): Naming template for spine switches.
                              The placeholder {} will be filled with a 1-based ID.
        output_filename (str): Name of the CSV file to generate.
        server_start_port (int): The port number on servers (1-based) to use for their connection.
                                 (Note: This is less relevant for the multi-port NVIDIA style but kept).
        server_ports_to_use (int): Number of ports on each server dedicated to
                                   leaf connections in NVIDIA style. Defaults to 8.
        server_port_prefix (str): Prefix for server port names (e.g., "eth") in NVIDIA style.
        leaf_port_prefix (str): Prefix for leaf port names (e.g., "Ethernet") in NVIDIA style.

    Returns:
        list: A list of tuples, where each tuple represents a connection
              (adev, aport, zdev, zport). This list is also written to the CSV.

    Raises:
        ValueError: If configurations are invalid (e.g., num_servers <=0,
                    odd ports_per_leaf, insufficient ports on switches for the topology).
    """
    connections = []

    if num_servers <= 0:
        # As per instruction, generate CSV. If num_servers is 0, it will be empty with headers.
        print("Number of servers is 0 or negative. Generating an empty CSV with headers.")
        with open(output_filename, 'w', newline='') as f:
            writer = csv.writer(f)
            writer.writerow(["adev", "aport", "zdev", "zport"])
        print(f"Generated empty '{output_filename}' with headers.")
        return connections


    if ports_per_leaf % 2 != 0:
        raise ValueError(
            f"ports_per_leaf ({ports_per_leaf}) must be an even number "
            "for the specified 50/50 downlink/uplink split."
        )
    
    downlink_ports_per_leaf = ports_per_leaf // 2
    uplink_ports_per_leaf = ports_per_leaf // 2 # Since ports_per_leaf is guaranteed even here

    if downlink_ports_per_leaf == 0: # Implies ports_per_leaf < 2
        raise ValueError(
            f"Leaf switches must have at least 2 ports to support downlink connections "
            f"(got ports_per_leaf = {ports_per_leaf}, resulting in {downlink_ports_per_leaf} downlink ports)."
        )

    # Calculate the number of leaf switches needed based on the total required downlink ports
    # for the NVIDIA style connections (num_servers * server_ports_to_use)
    num_leaf_switches = math.ceil((num_servers * server_ports_to_use) / downlink_ports_per_leaf)
    print(f"Calculated configuration: {num_leaf_switches} leaf switch(es).")

    # 1. Server-to-Leaf Connections (using NVIDIA style multi-port connections)
    server_to_leaf_conns = _generate_server_to_leaf_connections_nvidia_style(
        num_servers=num_servers,
        num_leaf_switches=num_leaf_switches,
        ports_per_leaf=ports_per_leaf, # Pass total ports per leaf
        server_template=server_template,
        leaf_template=leaf_template,
        server_ports_to_use=server_ports_to_use,
        server_port_prefix=server_port_prefix,
        leaf_port_prefix=leaf_port_prefix
    )
    connections.extend(server_to_leaf_conns)

    # 2. Leaf-to-Spine Connections
    if num_leaf_switches == 0: 
        # This case should ideally not be reached if num_servers > 0 and downlink_ports_per_leaf > 0.
        # If it were, no leaves means no leaf-spine connections.
        pass 
    elif uplink_ports_per_leaf == 0:
        # Leaves exist but have no capacity for uplinks.
        if num_leaf_switches > 1:
            print("Warning: Leaf switches have no uplink ports. Multiple leaves will exist but cannot be interconnected via a spine layer.")
        else: # num_leaf_switches == 1
            print("Note: A single leaf switch exists and has no uplink ports. No spine layer is needed or possible.")
    else:
        # Leaves exist and have uplink capacity. Proceed to define spine layer.
        # In a typical fat-tree, the number of spine switches is related to the uplink capacity of leaves.
        # Here, we assume each uplink port on a leaf connects to a distinct spine switch.
        num_spine_switches = uplink_ports_per_leaf
        print(f"Calculated configuration: {num_spine_switches} spine switch(es).")

        if num_spine_switches > 0:
            # Each spine switch must be able to connect to all leaf switches.
            # In this topology, spine port `k` (1-based) on any given spine connects to leaf `k`.
            # So, a spine needs at least `num_leaf_switches` ports.
            if ports_per_spine < num_leaf_switches:
                raise ValueError(
                    f"Spine switches have {ports_per_spine} ports, but need at least "
                    f"{num_leaf_switches} ports to connect to all {num_leaf_switches} leaf switches "
                    "(one port per leaf connected to that spine)."
                )

            for l_idx_0based in range(num_leaf_switches): # Iterate through each leaf switch
                current_leaf_name = leaf_template.format(l_idx_0based + 1)
                
                for u_on_leaf_idx_0based in range(uplink_ports_per_leaf): # Iterate through each uplink port on the current leaf
                    # Determine the actual port number on the leaf for this uplink (1-based)
                    # Uplink ports start after all downlink ports.
                    leaf_uplink_port_1based = downlink_ports_per_leaf + u_on_leaf_idx_0based + 1
                    
                    # Determine which spine this uplink connects to (0-indexed spine ID)
                    # The leaf's k-th uplink port connects to the k-th spine switch.
                    spine_idx_0based = u_on_leaf_idx_0based 
                    target_spine_name = spine_template.format(spine_idx_0based + 1)
                    
                    # Determine the port on the target spine switch (1-based)
                    # Spine `s` uses its port `l+1` to connect to leaf `l` (0-indexed l).
                    port_on_spine_1based = l_idx_0based + 1
                    
                    connections.append((
                        current_leaf_name,
                        leaf_uplink_port_1based, # Leaf's uplink port (adev, aport)
                        target_spine_name,
                        port_on_spine_1based     # Spine's port (zdev, zport)
                    ))
    
    # Write connections to CSV file
    with open(output_filename, 'w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["adev", "aport", "zdev", "zport"]) # Header
        for conn in connections:
            writer.writerow(conn)
    
    print(f"Successfully generated {len(connections)} connections in '{output_filename}'.")
    return connections


def main():
    parser = argparse.ArgumentParser(
        description="Generate a CSV file for network connections in a rail-optimized (2-level fat-tree) topology. "
                    "The CSV includes columns: adev, aport, zdev, zport.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter # Shows default values in help
    )
    
    parser.add_argument(
        "num_servers", 
        type=int, 
        help="Total number of servers (must be > 0)."
    )
    parser.add_argument(
        "--ports_per_leaf", 
        type=int, 
        default=64, 
        help="Number of ports per leaf switch. Must be an even number for the 50/50 downlink/uplink split."
    )
    parser.add_argument(
        "--ports_per_spine", 
        type=int, 
        default=64, 
        help="Number of ports per spine switch."
    )
    # Example for future extension to 3-level topology (Core layer)
    # parser.add_argument("--ports_per_core", type=int, default=32, help="Number of ports per core switch.")
    
    parser.add_argument(
        "--server_template", 
        type=str, 
        default="server{}", 
        help="Naming template for servers (e.g., 'gpu{}'). Use {} as placeholder for server number."
    )
    parser.add_argument(
        "--leaf_template", 
        type=str, 
        default="leaf{}", 
        help="Naming template for leaf switches. Use {} as placeholder for leaf switch number."
    )
    parser.add_argument(
        "--spine_template", 
        type=str, 
        default="spine{}", 
        help="Naming template for spine switches. Use {} as placeholder for spine switch number."
    )
    # parser.add_argument("--core_template", type=str, default="core{}", help="Naming template for core switches.")

    parser.add_argument(
        "--output_filename", 
        type=str, 
        default="cables.csv", 
        help="Name for the output CSV file."
    )
    parser.add_argument(
        "--server_start_port", 
        type=int, 
        default=1, 
        help="The starting port number (1-based) on servers for connections (less relevant for multi-port NVIDIA style)."
    )
    parser.add_argument(
        "--server_ports_to_use",
        type=int,
        default=8,
        help="Number of ports on each server to connect to leaf switches (NVIDIA style)."
    )
    parser.add_argument(
        "--server_port_prefix",
        type=str,
        default="",
        help="Prefix for server port names (e.g., 'eth') for NVIDIA style connections."
    )
    parser.add_argument(
        "--leaf_port_prefix",
        type=str,
        default="",
        help="Prefix for leaf port names (e.g., 'Ethernet') for NVIDIA style connections."
    )

    args = parser.parse_args()

    try:
        # Note: core_template and ports_per_core arguments are parsed if uncommented above,
        # but the generate_connections function currently does not use them.
        # This structure allows for future extension to a 3-level topology.
        generate_connections(
            num_servers=args.num_servers,
            ports_per_leaf=args.ports_per_leaf,
            ports_per_spine=args.ports_per_spine,
            server_template=args.server_template,
            leaf_template=args.leaf_template,
            spine_template=args.spine_template,
            output_filename=args.output_filename,
            server_start_port=args.server_start_port,
            server_ports_to_use=args.server_ports_to_use, # Pass new argument
            server_port_prefix=args.server_port_prefix,   # Pass new argument
            leaf_port_prefix=args.leaf_port_prefix,     # Pass new argument
        )
    except ValueError as e:
        print(f"Configuration Error: {e}")
    except Exception as e: # Catch any other unexpected errors
        print(f"An unexpected error occurred: {e}")

if __name__ == "__main__":
    main()
